from math import sqrt, ceil, floor
from gauseKernel import *
import numpy as np 

def edgeDetectKernels(kernel_type, sigma = 1):
    """
    2D Gause Kernel.
    """
    SUM = 0.0
    xtmp = [[1,0,-1],
            [2,0,-2],
            [1,0,-1]]
    ytmp = [[1,2,1],
            [0,0,0],
            [-1,-2,-1]]
    xKernel = np.array(xtmp)
    yKernel = np.array(ytmp)
    center = 1
    if kernel_type == "gause derivative":
        for i in range(3):
            for j in range(3):
                x = i - center
                y = j - center
                xKernel[i][j] = (-x / sigma**2) * math.exp( - (x**2 + y**2) / (2 * sigma**2))
                yKernel[i][j] = (-y / sigma**2) * math.exp( - (x**2 + y**2) / (2 * sigma**2))
    elif kernel_type == "sobel" or kernel_type == "Sobel":
        return xKernel, yKernel
    else:
        raise Exception("kernelTypeError, please enter 'gause derivative' or 'sobel'.")
    return xKernel, yKernel

class Conv2DWithEdgeDetection(Conv2D):
    def __init__(self, kernel_type, sigma = 1, boundary = "copy", padding = "SAME"):
        super().__init__(1, 3, sigma, boundary=boundary, padding=padding)
        xKernel, yKernel = edgeDetectKernels(kernel_type, sigma)
        self.filters = [xKernel, yKernel]

class CannyEdgeDetector():
    def __init__(self, image, kernel_type = "gause derivative", sigma = 1, boundary = "copy", padding = "SAME"):
        self.conv = Conv2DWithEdgeDetection(kernel_type, sigma, boundary, padding)
        self.image = image
    def getGrads(self):
        inputs = np.array(self.image)
        afterFliter = self.conv(inputs)
        xGrad = afterFliter[0]
        yGrad = afterFliter[1]
        self.xGrad = xGrad
        self.yGrad = yGrad
        return xGrad, yGrad
    def getMagnitudeAndOrientation(self):
        xGrad, yGrad = self.xGrad, self.yGrad
        h = len(xGrad)
        w = len(xGrad[0])
        self.h = h
        self.w = w
        magnitude = np.zeros(shape = (h, w))
        orientation = np.zeros(shape = (h, w, 2))
        for i in range(h):
            for j in range(w):
                magnitude[i][j] = sqrt(xGrad[i][j]**2 + yGrad[i][j]**2) + 0.000000001
                orientation[i][j][0] = yGrad[i][j] / magnitude[i][j]
                orientation[i][j][1] = xGrad[i][j] / magnitude[i][j]
        self.magnitude = magnitude
        self.orientation = orientation
        MagnitudeImage = Image.fromarray(self.magnitude.astype("uint")).convert("L")
        MagnitudeImage.save("./EdgeDetect/magnitude.jpg")
        return magnitude, orientation
    def nonMaxSuppression(self):
        h = len(self.magnitude)
        w = len(self.magnitude[0])
        for i in range(h):
            for j in range(w):
                rx = i + self.orientation[i][j][0]
                ry = j + self.orientation[i][j][1]
                px = i - self.orientation[i][j][0]
                py = j - self.orientation[i][j][1]
                rMagnitude = self.computeInterpolate(rx, ry, "near")
                pMagnitude = self.computeInterpolate(px, py, "near")
                if self.magnitude[i][j] < rMagnitude or self.magnitude[i][j] < pMagnitude:
                    self.magnitude[i][j] = 0
        NMSImage = Image.fromarray(self.magnitude.astype("uint")).convert("L")
        NMSImage.save("./EdgeDetect/NMS.jpg")
        return self.magnitude

    def computeInterpolate(self, x, y, type):
        if x < 0 or x >= self.h-1 or y < 0 or y >= self.w-1:
            return 0
        leftUpx = floor(x)
        leftUpy = floor(y)
        rightUpx = floor(x)
        rightUpy = ceil(y)
        leftDownx = ceil(x)
        leftDowny = floor(y)
        rightDownx = ceil(x)
        rightDowny = ceil(y)
        if type == "bilinear":
            magnitude = (leftUpx - x) * (leftUpy - y) * self.magnitude[rightDownx][rightDowny] + \
                        (rightDownx - x) * (rightDowny - y) * self.magnitude[leftUpx][leftUpy] + \
                        (rightUpx - x) * (rightUpy - y) * self.magnitude[leftDownx][leftDowny] + \
                        (leftDownx - x) * (leftDowny - y) * self.magnitude[rightUpx][rightUpy]
        elif type == "near":
            magnitude = max(self.magnitude[rightDownx][rightDowny], max(self.magnitude[leftUpx][leftUpy], max(self.magnitude[leftDownx][leftDowny], self.magnitude[rightUpx][rightUpy])))
        return magnitude
    def Thresholding(self, threshold):
        tmpMagnitude = np.array(self.magnitude)
        for i in range(self.h):
            for j in range(self.w):
                if tmpMagnitude[i][j] < threshold:
                    tmpMagnitude[i][j] = 0

        ThreImage = Image.fromarray(tmpMagnitude.astype("uint")).convert("L")
        ThreImage.save("./EdgeDetect/Thresholding-" + str(threshold) + ".jpg")
        return tmpMagnitude
    def hysteresisThresholding(self, high, low):
        tmpMagnitude = np.array(self.magnitude)
        for i in range(self.h):
            for j in range(self.w):
                if tmpMagnitude[i][j] < high:
                    tmpMagnitude[i][j] = 0
        
        
        for tmp in range(100):
            for i in range(self.h):
                for j in range(self.w):
                    if tmpMagnitude[i][j] != 0:
                        for m in range(-1,2):
                            for n in range(-1,2):
                                if i + m >=0 and i + m < self.h and j + n >= 0 and j + n < self.w:
                                    if self.magnitude[i+m][j+n] >= low and self.magnitude[i+m][j+n] < high:
                                        tmpMagnitude[i+m][j+n] = self.magnitude[i+m][j+n]
                                        
        hysThreImage = Image.fromarray(tmpMagnitude.astype("uint")).convert("L")
        hysThreImage.save("./EdgeDetect/hysteresisThresholding-" + str(high) + "-" + str(low) + ".jpg")
        return tmpMagnitude
    def run(self, high, low):
        self.getGrads()
        self.getMagnitudeAndOrientation()
        self.nonMaxSuppression()
        self.Thresholding(high)
        self.hysteresisThresholding(high, low)
imageRGB = Image.open("1.jpg")
image1 = imageRGB.convert("L")
edgeDetect = CannyEdgeDetector(image1, kernel_type="sobel")
edgeDetect.run(225,50)

